"""
sklearn是学习人工智能技术技术必须掌握的库，其中有一份糖尿病患者的数据sklearn.datasets中的load_diabetes，
请按照以下要求，正确实现线性回归：
"""
# 1)	导入相关的库，包括但不限于sklearn，numpy等等
import sklearn as sk
import numpy as np
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, optimizers, losses, metrics, callbacks
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score
from sklearn.preprocessing import StandardScaler

# 2)	从sklearn库中加载相应的数据
data = load_diabetes()

# 3)	打印输出数据中的键
print(data.keys())

# 4)	正确取出data和target中的数据分别作为X和Y
X = data['data']
Y = data['target']

# 5)	打印输出X，Y数据维度，并打印出前5个X,Y数据
print('X', X.shape)
print('Y', Y.shape)
print(X[:5])
print(Y[:5])

# 6)	仅取出X数据中的一个特征做训练
X = X[:, 2:3]
Y = Y.reshape(-1, 1)
X = StandardScaler().fit_transform(X)
Y = StandardScaler().fit_transform(Y)

# 7)	切分训练集，测试集
x_train, x_test, y_train, y_test = train_test_split(X, Y, train_size=0.7, random_state=1, shuffle=True)


# 8)	编写一个类，需要包含：代价函数，训练函数，测试函数
class DiabetesModel(keras.Model):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.fc = layers.Dense(1)

    def my_loss(self, alpha=0.1):
        # 9)	代价函数使用MSE均方误差
        self.compile(
            loss=losses.mean_squared_error,
            optimizer=optimizers.Adam(learning_rate=alpha),
        )

    def call(self, x, training=None):
        x = self.fc(x)
        return x

    # 10)	正确在类中编写训练的函数，设置迭代次数为20000，学习率为0.1
    def train(self, x, y):
        iters = 100
        batch_size = len(x)
        his = self.fit(x, y, epochs=iters,
                       batch_size=batch_size,
                       )
        return his

    # 11)	在类中编写出预测的方法
    def predict(self, x):
        return super().predict(x)


# 12)	实例化该类，并进行训练
model = DiabetesModel()
model.build((None, 1))
model.summary()
model.my_loss()
his = model.train(x_train, y_train)

# 13)	每500次打印，并输出代价
# 14)	绘出训练的学习曲线
loss_his = his.history['loss']
spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Loss')
plt.plot(loss_his)
plt.grid()

# 15)	绘出模型预测值和真实值的散点图
pred_test = model.predict(x_test)
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Pred and target')
plt.scatter(x_test, y_test, label='target')
plt.scatter(x_test, pred_test, label='prediction')
plt.grid()
plt.legend()
print('Please check and close the plotting window to continue ...')
plt.show()

# 16)	进行模型评估，计算R2值
r2 = r2_score(y_test, pred_test)
print(f'R2 = {r2}')

print('Over')
